#include "phy_simulator/visualizer.h"

#include <tf2_ros/transform_broadcaster.h>
#include <visualization_msgs/msg/marker_array.hpp>
#include <geometry_msgs/msg/point.hpp>

namespace phy_simulator {

Visualizer::Visualizer(std::shared_ptr<rclcpp::Node> node) : node_(node) {
  vehicle_set_pub_ = node_->create_publisher<visualization_msgs::msg::MarkerArray>(
      "vis/vehicle_set_vis", 10);
  lane_net_pub_ = node_->create_publisher<visualization_msgs::msg::MarkerArray>(
      "vis/lane_net_vis", 10);
  obstacle_set_pub_ = node_->create_publisher<visualization_msgs::msg::MarkerArray>(
      "vis/obstacle_set_vis", 10);
}

void Visualizer::VisualizeData() {
  auto time_stamp = node_->get_clock()->now();
  VisualizeDataWithStamp(time_stamp);
  // SendTfWithStamp(time_stamp);
}

void Visualizer::VisualizeDataWithStamp(const rclcpp::Time &stamp) {
  if (p_phy_sim_) {
    VisualizeVehicleSet(stamp, p_phy_sim_->vehicle_set());
    VisualizeLaneNet(stamp, p_phy_sim_->lane_net());
    VisualizeObstacleSet(stamp, p_phy_sim_->obstacle_set());
  } else {
    // If no physics simulation, publish empty markers
    visualization_msgs::msg::MarkerArray empty_marker;
    vehicle_set_pub_->publish(empty_marker);
    lane_net_pub_->publish(empty_marker);
    obstacle_set_pub_->publish(empty_marker);
  }
}

void Visualizer::SendTfWithStamp(const rclcpp::Time &stamp) {
  // TF broadcasting for ROS2 - can be implemented if needed
  // This would require tf2_ros::TransformBroadcaster
}

void Visualizer::VisualizeVehicleSet(const rclcpp::Time &stamp,
                                     const common::VehicleSet &vehicle_set) {
  visualization_msgs::msg::MarkerArray vehicle_marker;
  int marker_id = 0;
  
  for (auto iter = vehicle_set.vehicles.begin();
       iter != vehicle_set.vehicles.end(); ++iter) {
    const auto& vehicle = iter->second;
    
    // Create a simple cube marker for each vehicle
    visualization_msgs::msg::Marker marker;
    marker.header.frame_id = "map";
    marker.header.stamp = stamp;
    marker.ns = "vehicles";
    marker.id = marker_id++;
    marker.type = visualization_msgs::msg::Marker::CUBE;
    marker.action = visualization_msgs::msg::Marker::ADD;
    
    // Position
    marker.pose.position.x = vehicle.state().vec_position[0];
    marker.pose.position.y = vehicle.state().vec_position[1];
    marker.pose.position.z = 0.5;
    
    // Orientation
    marker.pose.orientation.x = 0.0;
    marker.pose.orientation.y = 0.0;
    marker.pose.orientation.z = sin(vehicle.state().angle / 2.0);
    marker.pose.orientation.w = cos(vehicle.state().angle / 2.0);
    
    // Scale
    marker.scale.x = vehicle.param().length();
    marker.scale.y = vehicle.param().width();
    marker.scale.z = 1.0;
    
    // Color
    marker.color.r = 0.8;
    marker.color.g = 0.8;
    marker.color.b = 0.8;
    marker.color.a = 1.0;
    
    vehicle_marker.markers.push_back(marker);
  }
  
  vehicle_set_pub_->publish(vehicle_marker);
}

void Visualizer::VisualizeLaneNet(const rclcpp::Time &stamp,
                                  const common::LaneNet &lane_net) {
  visualization_msgs::msg::MarkerArray lane_net_marker;
  int marker_id = 0;
  
  for (auto iter = lane_net.lane_set.begin(); iter != lane_net.lane_set.end();
       ++iter) {
    const auto& lane = iter->second;
    
    // Create line strip marker for the lane
    visualization_msgs::msg::Marker lane_marker;
    lane_marker.header.frame_id = "map";
    lane_marker.header.stamp = stamp;
    lane_marker.ns = "lanes";
    lane_marker.id = marker_id++;
    lane_marker.type = visualization_msgs::msg::Marker::LINE_STRIP;
    lane_marker.action = visualization_msgs::msg::Marker::ADD;
    
    // Scale
    lane_marker.scale.x = 0.1;
    lane_marker.scale.y = 0.1;
    lane_marker.scale.z = 0.1;
    
    // Color
    lane_marker.color.r = 0.0;
    lane_marker.color.g = 0.7;
    lane_marker.color.b = 1.0;
    lane_marker.color.a = 1.0;
    
    // Add points
    for (const auto& pt : lane.lane_points) {
      geometry_msgs::msg::Point point;
      point.x = pt(0);
      point.y = pt(1);
      point.z = 0.0;
      lane_marker.points.push_back(point);
    }
    
    lane_net_marker.markers.push_back(lane_marker);
  }
  
  lane_net_pub_->publish(lane_net_marker);
}

void Visualizer::VisualizeObstacleSet(const rclcpp::Time &stamp,
                                      const common::ObstacleSet &obstacle_set) {
  visualization_msgs::msg::MarkerArray obstacles_marker;
  int marker_id = 0;
  
  // Visualize circle obstacles
  for (const auto& circle_pair : obstacle_set.obs_circle) {
    const auto& obstacle = circle_pair.second;
    
    visualization_msgs::msg::Marker marker;
    marker.header.frame_id = "map";
    marker.header.stamp = stamp;
    marker.ns = "circle_obstacles";
    marker.id = marker_id++;
    marker.type = visualization_msgs::msg::Marker::CYLINDER;
    marker.action = visualization_msgs::msg::Marker::ADD;
    
    // Position
    marker.pose.position.x = obstacle.circle.center.x;
    marker.pose.position.y = obstacle.circle.center.y;
    marker.pose.position.z = 0.5;
    
    // Scale
    marker.scale.x = obstacle.circle.radius * 2;
    marker.scale.y = obstacle.circle.radius * 2;
    marker.scale.z = 1.0;
    
    // Color
    marker.color.r = 1.0;
    marker.color.g = 0.0;
    marker.color.b = 0.0;
    marker.color.a = 0.7;
    
    obstacles_marker.markers.push_back(marker);
  }
  
  // Visualize polygon obstacles
  for (const auto& polygon_pair : obstacle_set.obs_polygon) {
    const auto& obstacle = polygon_pair.second;
    
    visualization_msgs::msg::Marker marker;
    marker.header.frame_id = "map";
    marker.header.stamp = stamp;
    marker.ns = "polygon_obstacles";
    marker.id = marker_id++;
    marker.type = visualization_msgs::msg::Marker::LINE_STRIP;
    marker.action = visualization_msgs::msg::Marker::ADD;
    
    // Scale
    marker.scale.x = 0.1;
    
    // Color
    marker.color.r = 1.0;
    marker.color.g = 0.5;
    marker.color.b = 0.0;
    marker.color.a = 1.0;
    
    // Add points
    for (const auto& vertex : obstacle.polygon.points) {
      geometry_msgs::msg::Point point;
      point.x = vertex.x;
      point.y = vertex.y;
      point.z = 0.0;
      marker.points.push_back(point);
    }
    
    // Close the polygon
    if (!obstacle.polygon.points.empty()) {
      geometry_msgs::msg::Point point;
      point.x = obstacle.polygon.points[0].x;
      point.y = obstacle.polygon.points[0].y;
      point.z = 0.0;
      marker.points.push_back(point);
    }
    
    obstacles_marker.markers.push_back(marker);
  }
  
  obstacle_set_pub_->publish(obstacles_marker);
}

}  // namespace phy_simulator
